{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "d:\\cx\\Anaconda\\envs\\CX\\Lib\\site-packages\\torchrl\\data\\replay_buffers\\samplers.py:37: UserWarning: Failed to import torchrl C++ binaries. Some modules (eg, prioritized replay buffers) may not work with your installation. This is likely due to a discrepancy between your package version and the PyTorch version. Make sure both are compatible. Usually, torchrl majors follow the pytorch majors within a few days around the release. For instance, TorchRL 0.5 requires PyTorch 2.4.0, and TorchRL 0.6 requires PyTorch 2.5.0.\n",
      "  warnings.warn(EXTENSION_WARNING)\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "from torch import multiprocessing\n",
    "from collections import defaultdict\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from tensordict.nn import TensorDictModule\n",
    "from tensordict.nn.distributions import NormalParamExtractor\n",
    "from torch import nn\n",
    "from torchrl.collectors import SyncDataCollector\n",
    "from torchrl.data.replay_buffers import ReplayBuffer\n",
    "from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement\n",
    "from torchrl.data.replay_buffers.storages import LazyTensorStorage\n",
    "from torchrl.envs import (\n",
    "    Compose,\n",
    "    DoubleToFloat,\n",
    "    VecNorm,\n",
    "    StepCounter,\n",
    "    TransformedEnv,\n",
    ")\n",
    "from torchrl.envs.libs.gym import GymEnv\n",
    "from torchrl.envs.utils import check_env_specs, ExplorationType, set_exploration_type\n",
    "from torchrl.modules import ProbabilisticActor, TanhNormal, ValueOperator\n",
    "from torchrl.objectives import ClipPPOLoss\n",
    "from torchrl.objectives.value import GAE\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import cv2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "warnings.filterwarnings(\"ignore\")\n",
    "is_fork = multiprocessing.get_start_method() == \"fork\"\n",
    "device = (\n",
    "    torch.device(0)\n",
    "    if torch.cuda.is_available() and not is_fork\n",
    "    else torch.device(\"cpu\")\n",
    ")\n",
    "num_cells = 512\n",
    "lr = 1e-4\n",
    "max_grad_norm = 0.1\n",
    "frames_per_batch = 2048\n",
    "total_frames = 5_000_000\n",
    "sub_batch_size = 1024\n",
    "num_epochs = 10\n",
    "clip_epsilon = 0.2\n",
    "gamma = 0.99\n",
    "lmbda = 0.95\n",
    "entropy_eps = 1e-3\n",
    "save_path = \"./ppo_record2/\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_env = GymEnv(\n",
    "    \"Humanoid-v4\",\n",
    "    from_pixels=True,\n",
    "    pixels_only=False,\n",
    "    device=device,\n",
    ")\n",
    "env = TransformedEnv(\n",
    "    base_env,\n",
    "    Compose(\n",
    "        VecNorm(in_keys=[\"observation\", \"reward\"]),\n",
    "        DoubleToFloat(),\n",
    "        StepCounter(),\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running policy: TensorDict(\n",
      "    fields={\n",
      "        action: Tensor(shape=torch.Size([17]), device=cuda:0, dtype=torch.float32, is_shared=True),\n",
      "        done: Tensor(shape=torch.Size([1]), device=cuda:0, dtype=torch.bool, is_shared=True),\n",
      "        loc: Tensor(shape=torch.Size([17]), device=cuda:0, dtype=torch.float32, is_shared=True),\n",
      "        observation: Tensor(shape=torch.Size([376]), device=cuda:0, dtype=torch.float32, is_shared=True),\n",
      "        pixels: Tensor(shape=torch.Size([480, 480, 3]), device=cuda:0, dtype=torch.uint8, is_shared=True),\n",
      "        sample_log_prob: Tensor(shape=torch.Size([]), device=cuda:0, dtype=torch.float32, is_shared=True),\n",
      "        scale: Tensor(shape=torch.Size([17]), device=cuda:0, dtype=torch.float32, is_shared=True),\n",
      "        step_count: Tensor(shape=torch.Size([1]), device=cuda:0, dtype=torch.int64, is_shared=True),\n",
      "        terminated: Tensor(shape=torch.Size([1]), device=cuda:0, dtype=torch.bool, is_shared=True),\n",
      "        truncated: Tensor(shape=torch.Size([1]), device=cuda:0, dtype=torch.bool, is_shared=True)},\n",
      "    batch_size=torch.Size([]),\n",
      "    device=cuda:0,\n",
      "    is_shared=True)\n",
      "Running value: TensorDict(\n",
      "    fields={\n",
      "        done: Tensor(shape=torch.Size([1]), device=cuda:0, dtype=torch.bool, is_shared=True),\n",
      "        observation: Tensor(shape=torch.Size([376]), device=cuda:0, dtype=torch.float32, is_shared=True),\n",
      "        pixels: Tensor(shape=torch.Size([480, 480, 3]), device=cuda:0, dtype=torch.uint8, is_shared=True),\n",
      "        state_value: Tensor(shape=torch.Size([1]), device=cuda:0, dtype=torch.float32, is_shared=True),\n",
      "        step_count: Tensor(shape=torch.Size([1]), device=cuda:0, dtype=torch.int64, is_shared=True),\n",
      "        terminated: Tensor(shape=torch.Size([1]), device=cuda:0, dtype=torch.bool, is_shared=True),\n",
      "        truncated: Tensor(shape=torch.Size([1]), device=cuda:0, dtype=torch.bool, is_shared=True)},\n",
      "    batch_size=torch.Size([]),\n",
      "    device=cuda:0,\n",
      "    is_shared=True)\n"
     ]
    }
   ],
   "source": [
    "actor_net = nn.Sequential(\n",
    "    nn.LazyLinear(num_cells, device=device),\n",
    "    nn.LeakyReLU(),\n",
    "    nn.LazyLinear(num_cells, device=device),\n",
    "    nn.LeakyReLU(),\n",
    "    nn.LazyLinear(2 * env.action_spec.shape[-1], device=device),\n",
    "    NormalParamExtractor(),\n",
    ")\n",
    "policy_module = TensorDictModule(\n",
    "    actor_net, in_keys=[\"observation\"], out_keys=[\"loc\", \"scale\"]\n",
    ")\n",
    "policy_module = ProbabilisticActor(\n",
    "    module=policy_module,\n",
    "    spec=env.action_spec,\n",
    "    in_keys=[\"loc\", \"scale\"],\n",
    "    distribution_class=TanhNormal,\n",
    "    distribution_kwargs={\n",
    "        \"low\": env.action_spec.space.low,\n",
    "        \"high\": env.action_spec.space.high,\n",
    "    },\n",
    "    return_log_prob=True,\n",
    ")\n",
    "value_net = nn.Sequential(\n",
    "    nn.LazyLinear(num_cells, device=device),\n",
    "    nn.LeakyReLU(),\n",
    "    nn.LazyLinear(num_cells, device=device),\n",
    "    nn.LeakyReLU(),\n",
    "    nn.LazyLinear(1, device=device),\n",
    ")\n",
    "value_module = ValueOperator(module=value_net, in_keys=[\"observation\"])\n",
    "print(\"Running policy:\", policy_module(env.reset()))\n",
    "print(\"Running value:\", value_module(env.reset()))\n",
    "loss_module = ClipPPOLoss(\n",
    "    actor_network=policy_module,\n",
    "    critic_network=value_module,\n",
    "    clip_epsilon=clip_epsilon,\n",
    "    entropy_bonus=bool(entropy_eps),\n",
    "    entropy_coef=entropy_eps,\n",
    "    critic_coef=1.0,\n",
    "    loss_critic_type=\"smooth_l1\",\n",
    ")\n",
    "optim = torch.optim.Adam(loss_module.parameters(), lr)\n",
    "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n",
    "    optim, total_frames // frames_per_batch, 0.0\n",
    ")\n",
    "if 0:\n",
    "    checkpoint = torch.load(save_path + f\"weight/model_weights_test_{10}.pth\")\n",
    "    policy_module.load_state_dict(checkpoint[\"policy_weights\"])\n",
    "    value_module.load_state_dict(checkpoint[\"value_weights\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "collector = SyncDataCollector(\n",
    "    env,\n",
    "    policy_module,\n",
    "    frames_per_batch=frames_per_batch,\n",
    "    total_frames=total_frames,\n",
    "    split_trajs=False,\n",
    "    device=device,\n",
    ")\n",
    "replay_buffer = ReplayBuffer(\n",
    "    storage=LazyTensorStorage(max_size=frames_per_batch),\n",
    "    sampler=SamplerWithoutReplacement(),\n",
    ")\n",
    "advantage_module = GAE(\n",
    "    gamma=gamma,\n",
    "    lmbda=lmbda,\n",
    "    value_network=value_module,\n",
    "    average_gae=True,\n",
    "    device=device,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/5000000 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval cumulative reward: 11.3172 (init: 44.4515), eval step-count: 45, average reward=0.1979 (init=0.0050), step count (max): 62, lr policy: 0.0001:   1%|          | 59392/5000000 [19:19<22:51:05, 60.06it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "advantage中发现NaN！请检查归一化逻辑。\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "eval cumulative reward: 11.3172 (init: 44.4515), eval step-count: 45, average reward=0.1979 (init=0.0050), step count (max): 62, lr policy: 0.0001:   1%|          | 59392/5000000 [19:30<22:51:05, 60.06it/s]"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "reward中发现NaN奖励！请检查环境交互逻辑。",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mRuntimeError\u001b[39m                              Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[6]\u001b[39m\u001b[32m, line 9\u001b[39m\n\u001b[32m      7\u001b[39m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[33m\"\u001b[39m\u001b[33monservation中发现NaN！请检查归一化逻辑。\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m      8\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m tensordict_data[\u001b[33m\"\u001b[39m\u001b[33mnext\u001b[39m\u001b[33m\"\u001b[39m, \u001b[33m\"\u001b[39m\u001b[33mreward\u001b[39m\u001b[33m\"\u001b[39m].isnan().any():\n\u001b[32m----> \u001b[39m\u001b[32m9\u001b[39m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[33m\"\u001b[39m\u001b[33mreward中发现NaN奖励！请检查环境交互逻辑。\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m     10\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(num_epochs):\n\u001b[32m     11\u001b[39m     advantage_module(tensordict_data)\n",
      "\u001b[31mRuntimeError\u001b[39m: reward中发现NaN奖励！请检查环境交互逻辑。"
     ]
    }
   ],
   "source": [
    "base = 0\n",
    "logs = defaultdict(list)\n",
    "pbar = tqdm(total=total_frames)\n",
    "eval_str = \"\"\n",
    "for i, tensordict_data in enumerate(collector):\n",
    "    if tensordict_data[\"observation\"].isnan().any():\n",
    "        raise RuntimeError(\"onservation中发现NaN！请检查归一化逻辑。\")\n",
    "    if tensordict_data[\"next\", \"reward\"].isnan().any():\n",
    "        raise RuntimeError(\"reward中发现NaN奖励！请检查环境交互逻辑。\")\n",
    "    for _ in range(num_epochs):\n",
    "        advantage_module(tensordict_data)\n",
    "        if tensordict_data[\"advantage\"].isnan().any():\n",
    "            print(\"advantage中发现NaN！请检查归一化逻辑。\")\n",
    "            break\n",
    "        data_view = tensordict_data.reshape(-1)\n",
    "        replay_buffer.extend(data_view.cpu())\n",
    "        for _ in range(frames_per_batch // sub_batch_size):\n",
    "            subdata = replay_buffer.sample(sub_batch_size)\n",
    "            loss_vals = loss_module(subdata.to(device))\n",
    "            loss_value = (\n",
    "                loss_vals[\"loss_objective\"]\n",
    "                + loss_vals[\"loss_critic\"]\n",
    "                + loss_vals[\"loss_entropy\"]\n",
    "            )\n",
    "            loss_value.backward()\n",
    "            torch.nn.utils.clip_grad_norm_(loss_module.parameters(), max_grad_norm)\n",
    "            optim.step()\n",
    "            optim.zero_grad()\n",
    "    logs[\"reward\"].append(tensordict_data[\"next\", \"reward\"].mean().item())\n",
    "    pbar.update(tensordict_data.numel())\n",
    "    cum_reward_str = (\n",
    "        f\"average reward={logs['reward'][-1]:.4f} (init={logs['reward'][0]:.4f})\"\n",
    "    )\n",
    "    logs[\"step_count\"].append(tensordict_data[\"step_count\"].max().item())\n",
    "    stepcount_str = f\"step count (max): {logs['step_count'][-1]}\"\n",
    "    logs[\"lr\"].append(optim.param_groups[0][\"lr\"])\n",
    "    lr_str = f\"lr policy: {logs['lr'][-1]:.4f}\"\n",
    "    if i % 10 == 0:\n",
    "        env.eval()\n",
    "        with set_exploration_type(ExplorationType.DETERMINISTIC), torch.no_grad():\n",
    "            eval_rollout = env.rollout(1000, policy_module)\n",
    "            if eval_rollout[\"next\", \"reward\"].isnan().any():\n",
    "                raise RuntimeError(\"评估数据中发现NaN奖励！请检查环境或策略输出。\")\n",
    "            if eval_rollout[\"observation\"].isnan().any():\n",
    "                raise RuntimeError(\"评估观测中发现NaN！请检查归一化评估模式。\")\n",
    "            logs[\"eval reward\"].append(eval_rollout[\"next\", \"reward\"].mean().item())\n",
    "            logs[\"eval reward (sum)\"].append(\n",
    "                eval_rollout[\"next\", \"reward\"].sum().item()\n",
    "            )\n",
    "            logs[\"eval step_count\"].append(eval_rollout[\"step_count\"].max().item())\n",
    "            eval_str = (\n",
    "                f\"eval cumulative reward: {logs['eval reward (sum)'][-1]:.4f} \"\n",
    "                f\"(init: {logs['eval reward (sum)'][0]:.4f}), \"\n",
    "                f\"eval step-count: {logs['eval step_count'][-1]}\"\n",
    "            )\n",
    "\n",
    "            ims = eval_rollout[\"pixels\"].cpu().numpy().astype(\"uint8\")\n",
    "            fourcc = cv2.VideoWriter_fourcc(*\"mp4v\")\n",
    "            out = cv2.VideoWriter(\n",
    "                save_path + f\"video/output_test_{i + base}.mp4\",\n",
    "                fourcc,\n",
    "                30,\n",
    "                (ims.shape[2], ims.shape[1]),\n",
    "            )\n",
    "            for frame in ims:\n",
    "                out.write(frame[..., ::-1])\n",
    "            out.release()\n",
    "            del eval_rollout\n",
    "        env.train()\n",
    "        torch.save(\n",
    "            {\n",
    "                \"policy_weights\": policy_module.state_dict(),\n",
    "                \"value_weights\": value_module.state_dict(),\n",
    "            },\n",
    "            save_path + f\"weight/model_weights_test_{i + base}.pth\",\n",
    "        )\n",
    "    pbar.set_description(\", \".join([eval_str, cum_reward_str, stepcount_str, lr_str]))\n",
    "    scheduler.step()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
